#![allow(unsafe_code)]

use stdx::ptr::{cast_read_at, cast_write_at, read_at};
use zuc::ZUC128;

pub fn encryption(input: impl AsRef<[u8]>) -> String {
    let encrypted = zuc128_pkcs7_encrypt(input.as_ref());
    base64_simd::STANDARD.encode_to_string(&encrypted)
}

fn zuc128_pkcs7_encrypt(input: &[u8]) -> Vec<u8> {
    const B: usize = 4;
    let pad_len = B - input.len() % B;
    let ans_len = input.len() + pad_len;

    let mut ans = Vec::with_capacity(ans_len);

    let k = [0u8; 16];
    let iv = [1u8; 16];
    let mut zuc = ZUC128::new(&k, &iv);

    unsafe {
        let src: *const u8 = input.as_ptr();
        let dst: *mut u8 = ans.as_mut_ptr();
        for i in 0..input.len() / B {
            let block = u32::from_be_bytes(cast_read_at(src, i * B));
            let key = zuc.generate();
            let encrypted = block ^ key;
            cast_write_at(dst, i * B, encrypted.to_be_bytes());
        }

        let last = input.len() / B * B;
        let last_block = match input.len() % B {
            0 => 0x0404_0404,
            1 => u32::from_be_bytes([
                read_at(src, last), //
                0x03,
                0x03,
                0x03,
            ]),
            2 => u32::from_be_bytes([
                read_at(src, last), //
                read_at(src, last + 1),
                0x02,
                0x02,
            ]),
            3 => u32::from_be_bytes([
                read_at(src, last),
                read_at(src, last + 1),
                read_at(src, last + 2),
                0x01,
            ]),
            _ => unreachable!(),
        };
        let key = zuc.generate();
        let encrypted = last_block ^ key;

        cast_write_at(dst, last, encrypted.to_be_bytes());

        ans.set_len(ans_len);
        ans
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn debug() {
        let (input, expected) = ("特朗普", "/QZb7S0JWp/IYuwB");
        eprintln!("input: {:02x?}", input.as_bytes());

        let expected_bytes = base64_simd::STANDARD.decode_to_vec(expected).unwrap();
        let output_bytes = zuc128_pkcs7_encrypt(input.as_bytes());
        eprintln!("output:   {output_bytes:02x?}");
        eprintln!("expected: {expected_bytes:02x?}");
        assert_eq!(output_bytes, expected_bytes);

        let output = encryption(input);
        assert_eq!(output, expected);
    }
}
